# 网络地址转换：内网->外网, 端口号也变
from collections import defaultdict
from time import time_ns

from scapy.compat import raw
from scapy.layers.inet import IP, TCP, UDP


# Simple NAT, rewrite ip address and port
class NAT(object):
    def __init__(self, public_address: str, public_port: int = 8000):
        self.public_address = public_address
        # 端口+协议
        self.tcp_port = 1024  # 从1024开始映射
        self.udp_port = 1024
        # 指定映射端口
        self.public_port = public_port
        self.send_port_map = defaultdict(tuple)  # 记录映射的端口
        self.receive_port_map = defaultdict(tuple)

    # 传入时最好只有包头
    def match(self, pkt: bytes) -> bytes:  # 内网->外网| 外网-> 内网
        # 收取发送方的包
        scapy_pkt = IP(pkt)
        src, dst = scapy_pkt[IP].src, scapy_pkt[IP].dst
        # TCP 或 UDP
        proto = scapy_pkt[IP].proto
        if proto == 6:
            src_port = scapy_pkt[TCP].sport
            dst_port = scapy_pkt[TCP].dport
            flag = 'TCP'
        else:
            flag = 'UDP'
            src_port = scapy_pkt[UDP].sport
            dst_port = scapy_pkt[UDP].dport

        if src != self.public_address and dst != self.public_address:  # 内网->外网
            key = (flag, src, src_port)
            if self.public_port is not None:
                port = self.public_port
            elif flag == 'TCP':
                self.tcp_port += 1
                port = self.tcp_port
            else:
                self.udp_port += 1
                port = self.udp_port
            if key not in self.send_port_map:
                self.send_port_map[key] = (self.public_address, port)
                # 反过来记录
                self.receive_port_map[(flag, self.public_address, port)] = (src, src_port)
            scapy_pkt[IP].src = self.public_address
            if flag == 'TCP':
                scapy_pkt[TCP].sport = port
                del scapy_pkt[TCP].chksum
            else:
                scapy_pkt[UDP].sport = port
                del scapy_pkt[UDP].chksum
                del scapy_pkt[UDP].len
            del scapy_pkt[IP].chksum
            return bytes(IP(raw(scapy_pkt)))  # 自动计算

        if dst == self.public_address and dst_port > 1024:
            # 返回的包
            key = (flag, dst, dst_port)
            if key in self.receive_port_map:
                dst, dst_port = self.receive_port_map[key]

                scapy_pkt[IP].dst = dst
                if flag == 'TCP':
                    scapy_pkt[TCP].dport = dst_port
                    del scapy_pkt[TCP].chksum
                else:
                    scapy_pkt[UDP].dport = dst_port
                    del scapy_pkt[UDP].chksum
                    del scapy_pkt[UDP].len
            del scapy_pkt[IP].chksum
            return bytes(IP(raw(scapy_pkt)))  # 自动计算
        return pkt


if __name__ == '__main__':
    nat = NAT("172.17.0.2", 8000)
    pkt1 = b'E\x00\x00-\x00\x01\x00\x00@\x06\x11\xf4\xf7\x89\xe9\xcc$\x91b\xef\x8duM\xf1\x00\x00\x00\x00\x00\x00\x00' \
           b'\x00P\x02 \x00!\xd9\x00\x00'
    # IP(pkt1).show()
    pkt = nat.match(pkt1)
    # IP(pkt).show()
    pkt2 = b'E\x00\x00@\x00\x01\x00\x00@\x06\x11\xe1\xf7\x89\xe9\xcc$\x91b\xeffeM\xf1\x00\x00\x00\x00\x00\x00\x00\x00P' \
           b'\x02 \x00\xe2\x0c\x00\x00aaaaaaaaaaaaaaaaaaaaaaaa '
    pkt = IP(src='247.137.233.204', dst='36.145.98.239') / TCP(sport=26213, dport=19953)
    pkt[TCP].add_payload(b'a' * 2008)
    pkt2 = raw(pkt)
    # IP(pkt2).show()
    pkt = nat.match(pkt2)
    # IP(pkt).show()
    start = time_ns()
    for i in range(100000):
        # nat.translate(pkt1)
        nat.match(pkt2)
    print((time_ns() - start) / 100000 / 1000)  # 平均每条指令执行时间
